import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 导入模型
my_model = tf.keras.models.load_model("clothing_model.keras")
clothing_model = tf.keras.Sequential([my_model, tf.keras.layers.Softmax()])

# 导入测试数据集。
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

x_test = x_test / 255.0  # 归一化

# 预测测试集。
predictions = clothing_model.predict(x_test)

print(np.argmax(predictions[0]))

def main():
    num_rows = 5  # 定义一个行数
    num_cols = 3  # 定义一个列数
    num_images = num_rows * num_cols

    fig, axes = plt.subplots(num_rows, num_cols * 2, figsize=(2 * 2 * num_cols, 2 * num_rows))

    for i in range(num_images):
        pradict_data = clothing_model.predict(x_test[i][np.newaxis, ...]).squeeze()
        predict_label = np.argmax(pradict_data)

        row = i // num_cols
        col = i % num_cols

        ax = axes[row, col * 2]
        ax.set_xticks([])
        ax.set_yticks([])
        ax.grid(False)
        ax.imshow(x_test[i], cmap=plt.cm.binary)

        color = 'blue' if predict_label == y_test[i] else 'red'
        ax.set_xlabel(
            f'{class_names[predict_label]} {100*np.max(pradict_data):.2f}% ({class_names[y_test[i]]})',
            color=color
        )

        ax = axes[row, col * 2 + 1]
        ax.grid(False)
        ax.set_yticks([])
        ax.set_xticks(range(10))
        ax.set_ylim([0, 1])

        thisplot = ax.bar(range(10), pradict_data, color='#777777')
        thisplot[predict_label].set_color('red')
        thisplot[y_test[i]].set_color('blue')

    fig.tight_layout()
    plt.show()


if __name__ == '__main__':
    main()
